package edu.map.aamas.rnn;

import java.util.Random;
import java.util.Timer;
import java.util.TimerTask;
import java.util.Vector;

import edu.map.aamas.gui.Tadpole;

public class RNN {

	// dynamic values
	public double[][] M;
	public double[] X;
	public double[] H;
	public double[] P;
	public double[] N;
	public double[] C;
	
	// parameter values
	public int size;	
	public double xmin = -1;
	public double xmax = 1;
	public double _a1 = 0.1;
	public double _a2 = 0.1;
	private double x0 = 0;
	public int numMuscles = 2;
	
	public long iteration = 1;
	
	public String[] possibleControllers = {"RNR", "Harmonics", "Random"};
	public String currentController = "RNR";
	public int millisBetweenIterations = 50;
	
	public Vector<RnnListener> listeners = new Vector<RnnListener>();

	public long randomSeed = System.currentTimeMillis();
	public boolean useRandomSeed = true;
	public boolean paused = true;
	
	private Timer updateTimer;
	private TimerTask updaterTask;
	public static RNN instance = null;
	public boolean proprioception = true;
	public boolean homeostasis = true;
	public boolean plasticity = false;
	public boolean loggingEnabled = false;
	public static RNN getInstance() {
		if (instance == null)
			instance = new RNN(32);
		return instance;
	}
	
	public static RNN getInstance(int size) {
		if (instance == null)
			instance = new RNN(size);
		
		if (instance.size != size) {
			instance.pause();
			instance.size = size;
			instance.init();
			instance.resume();
		}
		return instance;
	}
	
	private RNN(int size) {		
		this.size = size;
		init();     	
	}
	
	public double kineticEnergy() {
		double sum = 0;
		for (int i = 0; i < size; i++)
			sum += X[i];
		
		return sum/size;
	}
	
	public synchronized void init() {
		
		M = new double[size][size];
		X = new double[size];
		H = new double[size];
		P = new double[size];
		C = new double[numMuscles];

		for (int i = 0; i < size; i++) {
			X[i] = x0;
			H[i] = 0;
			P[i] = 0;
		}
		Random r;
		
		if (useRandomSeed)
			randomSeed = System.currentTimeMillis();

		r = new Random(randomSeed);
		
		for (int i = 0; i < size; i++) {
			for (int j=0; j < size; j++) {				
				if (i != j)
					M[i][j] = r.nextGaussian();
			}
		}
		iteration = 1;
		
		for (RnnListener l : listeners)
			l.rnnRestarted(this);
	}

	
	public void update() {
		//System.out.println("update");
		if (currentController.equals("RNR")) {
			if (homeostasis)
				updateH();					
			if (proprioception)
				updateP();
			if (plasticity)
				updateM();
			updateX();			
			updateC();						
		}
		else if (currentController.equalsIgnoreCase("Harmonics")) {
			updateX_harmonics();
			updateC();			
		}
		else if (currentController.equalsIgnoreCase("Random")) {
			updateX_random();
			updateC();
		}
		
		iteration++;
		updateListeners();		
	}

	// Simple neural network update
	private void updateX() {
		for (int i=0; i<size;i++) {
			double act = 0;
			for (int j=0; j<size;j++) {
				if (j!=i) 
					act = act + X[j]*M[j][i];  
			}
			X[i] = X[i]*(1-_a1) + _a1*(act - H[i] + P[i]);
			X[i] = minmax(X[i], xmin, xmax);
		}
	}
	
	private double minmax(double val, double min, double max) {
		  if (val >= max) 
			  return max;
		  if (val <= min) 
			  return min;
		  return val;
	}
	
	// Homeostasis
	private void updateH() {
		for (int i = 0; i < size; i++) {
			H[i] = H[i]*(1-_a2) + _a2*(X[i]-x0);
		}
	}
	
	// Neural Plasticity
	public double alpha3 = 0.03;
	private void updateM() {
		for (int i = 0; i < size; i++)
			for (int j = 0; j < size; j++) {
				if (j == i)
					continue;
				
				M[i][j] = M[i][j] * (1-alpha3) + alpha3 * plasticityF(i, j);			
			}
	}
	
	private void updateP_old() {
		//FIXME
		for (int i = size%2; i < size - 1; i+=2){
			P[i] = Tadpole.getInstance().instantSpeed;
			P[i+1] = minmax(Math.abs(Tadpole.getInstance().angularSpeed*5), 0, 4);
		}
		
	}
	

	public double k1 = 0.2;
	public double k2 = 0.3;
	
	private void updateP() {
		double angDiff1 = (Tadpole.getInstance().getTailAngle() - Tadpole.FORCE_ANGLE_FACTOR);
		double angDiff2 = (Tadpole.getInstance().getTailAngle() + Tadpole.FORCE_ANGLE_FACTOR);
		
		angDiff1 = angDiff1 * angDiff1 * k2;
		angDiff2 = angDiff2 * angDiff2 * k2;
		
		double p1 = k1 * Math.exp(angDiff1);
		double p2 = k1 * Math.exp(angDiff1);
		
		//System.out.println(p1+", "+p2);
		for (int i = 0; i < size/2; i++){
			P[i] = p1;			
		}
		for (int i = size/2; i < size; i++){
			P[i] = p2;			
		}
		
	}
	
	
	private void updateC() {

		for (int i = 0; i < numMuscles; i++)
			C[i] = 0;
		
		int startIndex = size%numMuscles;
		int muscleNeurons = size/(numMuscles-startIndex);
		
		// if the neurons can't be exactly divided, ignore some of them 
		
		for (int i=startIndex; i<size; i++) {			
			C[i/muscleNeurons] += X[i] - x0;
		}
	}
	
	private void updateX_harmonics() {
		
		for (int i=0; i<size/2; i++) {			
			X[i] = 0.2 * Math.sin(0.4 * iteration);
		}
		
		for (int i=size/2; i<size; i++) {			
			X[i] = -0.2 * (Math.sin(0.4 * iteration));
		}
	}
	
private void updateX_random() {
		Random r = new Random(System.currentTimeMillis());
		for (int i=0; i<size; i++) {			
			X[i] = 0.1 * r.nextGaussian();
		}
	}
	
	private void printArray(double[] array) {
		for (int i = 0; i < array.length; i++)
			System.out.print(array[i]+" ");
		System.out.println();
	}

	public void printState() {
		System.out.println("\nX:");
		printArray(X);
		System.out.println("\nH:");
		printArray(H);
		System.out.println("\nP:");
		printArray(P);
		System.out.println("\nC:");
		printArray(C);
		System.out.println("force: "+getDiffC0C1());
	}
	
	public double getDiffC0C1() {
		double diff = C[0] - C[1];
		diff /= size;
		return diff;
	}
	
	public void addRnrListener (RnnListener listener) {
		if (!listeners.contains(listener))
			listeners.add(listener);
	}
	
	public void removeRnrListener(RnnListener listener) {
		listeners.remove(listener);
	}
	
	public void updateListeners() {
		for (RnnListener listener : listeners)
			listener.rnnUpdated(this);
	}
	
	public static void main(String[] args) {
		RNN rnr = new RNN(16);
		int i = 0;
		while (true) {
			rnr.printState();
			rnr.update();
			i++;
		}
	}
	
	public synchronized void pause() {	
		paused = true;
		if (updateTimer != null) {
			updateTimer.cancel();
			updateTimer = null;
		}
	}
	
	public synchronized void resume() {		
		if (paused) {
			paused = false;
			startUpdating();
		}
	}
	
	
	public synchronized void startUpdating() {

		if (updateTimer != null)
			updateTimer.cancel();
		
		updaterTask = new TimerTask() {
			public void run() {	
				while (!paused) {
					update();
					try {
						Thread.sleep(millisBetweenIterations);
					}
					catch (Exception e) {
						e.printStackTrace();
					}
				}			
			}
		};
		
		updateTimer = new Timer("RNR updater");
		updateTimer.schedule(updaterTask, millisBetweenIterations, millisBetweenIterations);
	}
	
	private double plasticityG(double a) {
		return 1 - (a/(xmax - xmin));
	}
	
	private double plasticityF(int i, int j) {
		double sign = Math.signum(M[i][j]);
		return sign * plasticityG(Math.abs(X[i])) * plasticityG(Math.abs(X[j]));
	}
}